import pandas as pd
import json
from pyecharts import options as opts
from pyecharts.charts import Tree

# 1. 读取广告转化数据集
# 数据源：https://www.kaggle.com/loveall/clicks-conversion-tracking/data
# 广告转化数据集：Ad Conversions Data
df = pd.read_csv('./Files/KAG_conversion_data.csv')
print(df.head())

# 数据处理
df["total"] = "全量"
df["xyz_campaign_id"] = df["xyz_campaign_id"].map(lambda x: f"c_{x}")
df["gender"] = df["gender"].map(lambda x: f"g_{x}")
df["age"] = df["age"].map(lambda x: f"a_{x}")

# 2. 准备数据：全量 > 活动 > 性别 > 年龄 四级树
df_agg = df.groupby(["total", "xyz_campaign_id", "gender", "age"]).size().rename("size").reset_index()
print(df_agg)

data = {}
data['name'] = "全量"
data['children'] = []
for campaign in df_agg['xyz_campaign_id'].unique():
    campaign_node = {}
    campaign_node['name'] = campaign
    campaign_node['children'] = []

    for gender in df_agg[df_agg['xyz_campaign_id'] == campaign]['gender'].unique():
        gender_node = {}
        gender_node['name'] = gender
        gender_node['children'] = []

        age_datalist = df_agg[(df_agg['xyz_campaign_id'] == campaign) & (df_agg['gender'] == gender)][['age', 'size']]

        for age, value in age_datalist.values:
            age_node = {}
            age_node['name'] = age
            age_node['value'] = value
            gender_node['children'].append(age_node)

        campaign_node['children'].append(gender_node)

    data['children'].append(campaign_node)

print(json.dumps(data))

tree = (
    Tree()
    .add(
        "",
        [data],
        # 参数可以有：取值分别为 'LR' , 'RL', 'TB', 'BT'
        orient="LR",
        # layout：orthogonal，radial，如果是radial则是圆环发散图
        # layout="radial",
        # 定义树状图的名称
        label_opts=opts.LabelOpts(
            position="top",
            horizontal_align="right",
            vertical_align="middle"
        )
    )
    .set_global_opts(title_opts=opts.TitleOpts(title="广告流量图", pos_left="center"))
)

tree.render('./Files/my_tree.html')
